package main

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"sort"
	"spark-port-proxy/util"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Server struct {
	Queue            chan ReceivePackage
	BroadcastAddress string
	Port             int
	HttpPort         int
	Interval         time.Duration
	state            string
	Hostname         string
	IP               string
	nodeId           int
	upTime           int64
	upNum            int
	online           *sync.Map
	ProxyStartPort   string // 代理的起始端口，支持配置多个，逗号分割
	ProxyPorts       int    //从起始端口开始，向后增加的端口个数
	listenPorts      map[int]net.Listener
	deadNode         map[int]Node
	allPortTable     map[int]string
	mutex            sync.Mutex
	listenMutex      sync.Mutex //监听端口锁，保证并发先map正确更新
}

type Node struct {
	ID       int    // 节点ID
	Time     int64  // 时间戳
	Health   string // 节点健康状态
	Port     int    // 分配的端口
	UpTime   int64
	Hostname string
}

const UNREADY_STATE = "unready"
const READY_STATE = "ready"

type ReceivePackage struct {
	IP      string
	Content string
}

func (receiver *Server) receiveUdp() {
	conn, err := net.ListenPacket("udp4", fmt.Sprintf(":%d", receiver.Port))
	if err != nil {
		fmt.Println("Failed to start UDP listener:", err)
		return
	}
	defer conn.Close()
	log.Printf("listen udp: %d \n", receiver.Port)

	buffer := make([]byte, 1024)

	for {
		n, addr, err := conn.ReadFrom(buffer)
		if err != nil {
			fmt.Println("Error while reading UDP message:", err)
			continue
		}
		senderIP := addr.(*net.UDPAddr).IP.String()
		message := string(buffer[:n])

		// log.Println("push message in queue: " + message)
		queue <- ReceivePackage{
			IP:      senderIP,
			Content: message,
		}
	}
}

// 开启服务
func (receiver *Server) Serve() {
	receiver.state = "unready"
	receiver.online = new(sync.Map)
	receiver.allPortTable = make(map[int]string)
	receiver.deadNode = make(map[int]Node)
	receiver.listenPorts = make(map[int]net.Listener)
	receiver.rollNodeId()
	receiver.updatePorts([]int{})
	go receiver.receiveUdp()
	go receiver.reply()
	go receiver.broadcast()
	receiver.serveHttp()
}

func (receiver *Server) broadcast() {
	broadcastTicker := time.NewTicker(receiver.Interval * time.Second)
	defer broadcastTicker.Stop()

	for {
		select {
		case <-broadcastTicker.C:
			if receiver.state == UNREADY_STATE && receiver.upNum < 3 {
				// 版本号：操作码：类型：序号：内容
				message := fmt.Sprintf("001:1:0:%s#%s:%d:%d:%s", receiver.IP, receiver.Hostname, receiver.nodeId, receiver.upTime, receiver.state)
				util.SendPackage(receiver.BroadcastAddress, receiver.Port, message)
				receiver.upNum++
			} else {
				receiver.state = READY_STATE
				// 当状态已经准备好，即已经确认ID，通知其他节点当前节点的状态
				message := fmt.Sprintf("001:1:0:%s#%s:%d:%d:%s", receiver.IP, receiver.Hostname, receiver.nodeId, receiver.upTime, receiver.state)
				util.SendPackage(receiver.BroadcastAddress, receiver.Port, message)
				receiver.upNum = 0
			}
			// 检查在线的节点是否存活着
			receiver.onlineAlive()
		}
	}
}

func (receiver *Server) reply() {
	for message := range receiver.Queue {
		index := strings.Index(message.Content, "#")
		if index < 8 {
			continue
		}
		header := message.Content[0:index]
		text := message.Content[index+1:]
		strs := strings.Split(header, ":")
		if len(strs) != 4 {
			continue
		}
		op := strs[1]
		inOut := strs[2]
		ip := strs[3]
		if strings.HasSuffix(message.IP, ".1") {
			message.IP = ip
		}
		// log.Println("receive message: " + message.IP)

		if strings.HasSuffix(text, UNREADY_STATE) {
			log.Println(message.Content)
		}
		// Server -> Client: Client
		if op == "1" && inOut == "0" {
			actInfo := strings.Split(text, ":")
			hostname := actInfo[0]
			id, _ := strconv.Atoi(actInfo[1])
			upTime, _ := strconv.ParseInt(actInfo[2], 10, 64)
			state := actInfo[3]
			// 如果ID重复，且启动时间晚，则答复需要更换ID
			if receiver.Hostname != hostname && receiver.nodeId == id && upTime >= receiver.upTime && state == UNREADY_STATE {
				replyMessage := fmt.Sprintf("001:1:1:%s#%d:0", receiver.IP, id)
				util.SendPackage(message.IP, receiver.Port, replyMessage)
				log.Println("node id equals with me:", id)
			}
			// 如果发送消息的节点已经是准备状态，则记录到活动列表中
			if state == READY_STATE {
				value, exist := receiver.online.Load(id)
				var node Node
				if !exist {
					node = Node{
						ID:       id,
						UpTime:   upTime,
						Hostname: hostname,
					}
				} else {
					node = value.(Node)
				}
				// online数量变更时，需要进入互斥锁，防止在遍历的时候数量不统一
				receiver.mutex.Lock()
				node.Time = time.Now().UnixNano()
				receiver.online.Store(id, node)
				value, _ = receiver.online.Load(id)
				// log.Println(value)
				if !exist {
					log.Println("new node join: " + hostname)
					receiver.onlineChange()
				}
				receiver.mutex.Unlock()
			}
		}
		// Client -> Server: Server
		if op == "1" && inOut == "1" {
			actInfo := strings.Split(text, ":")
			id, _ := strconv.Atoi(actInfo[0])
			code, _ := strconv.Atoi(actInfo[1])
			if id == receiver.nodeId && code == 0 {
				receiver.rollNodeId()
			}
		}

		// Nginx -> Proxy: request to query host by port
		// Never run code, because udp ignore response
		if op == "2" && inOut == "0" {
			port, err := strconv.Atoi(text)
			if err != nil {
				log.Println("request to query host failed: ", err)
				continue
			}
			host, exist := receiver.allPortTable[port]
			code := 0
			if exist {
				code = 1
			}
			replyMessage := fmt.Sprintf("001:2:1:10664#%d:%s", code, host)
			util.SendPackage(message.IP, receiver.Port, replyMessage)
		}
	}
}

func (receiver *Server) rollNodeId() {
	receiver.nodeId++
	receiver.upTime = time.Now().UnixNano()
	receiver.upNum = 0
	log.Println("roll node id:", receiver.nodeId)
}

/**
当在线节点列表发生变化时（新加入或者掉线退出），需要重新分配端口
*/
func (receiver *Server) onlineChange() {
	strPorts := strings.Split(receiver.ProxyStartPort, ",")
	ports := make([]int, len(strPorts))
	for i, value := range strPorts {
		ports[i], _ = strconv.Atoi(value)
	}
	online := util.ToMap(*receiver.online)
	log.Println("begin ite online key:", len(online))
	// 根据数量，排序节点后，从前往后分配端口，有些可能分配不到
	var keys []int
	for key := range online {
		keys = append(keys, key.(int))
		log.Println("ite online key:", key)
	}
	sort.Ints(keys)
	pointer := 0
	for index, key := range keys {
		if key == receiver.nodeId {
			pointer = index
			break
		}
	}
	log.Println(keys)
	var planPorts []int
	// node id start with 1, or not 0
	p := 0
	for i := 0; i < receiver.ProxyPorts; i++ {
		value, _ := online[keys[p]]
		log.Println("ite node:", keys[p])
		node := value.(Node)
		for _, port := range ports {
			receiver.allPortTable[port+i] = node.Hostname
			if p == pointer {
				planPorts = append(planPorts, port+i)
				log.Println(node.Hostname+" allocate port:", port+i)
			}
		}
		p++
		if p >= len(online) {
			p = 0
		}
	}
	receiver.updatePorts(planPorts)

}

func (receiver *Server) onlineAlive() {
	receiver.mutex.Lock()
	defer receiver.mutex.Unlock()

	online := util.ToMap(*receiver.online)
	now := time.Now().UnixNano()
	// 等待下次检查时，将未响应的节点移除
	size := 0
	for i := range receiver.deadNode {
		value, _ := online[i]
		node := value.(Node)
		//  再次确认最后一次响应时间
		if now-node.Time > (10 * time.Second).Nanoseconds() {
			delete(receiver.deadNode, i)
			delete(online, i)
			receiver.online.Delete(i)
			log.Println("current time:", now)
			log.Println("remove dead node:", node)
			size++
		} else {
			delete(receiver.deadNode, i)
		}
	}
	if size > 0 {
		receiver.onlineChange()
	}
	for i := range online {
		id := i.(int)
		value, _ := online[i]
		node := value.(Node)
		// 最后一次活跃时间不能大于10s，否则认为是掉线
		if now-node.Time > 10*1000*1000 {
			receiver.deadNode[id] = node
		}
	}
}

func (receiver *Server) updatePorts(ports []int) {
	strPorts := strings.Split(receiver.ProxyStartPort, ",")
	allPorts := make([]int, len(strPorts))
	for i, value := range strPorts {
		allPorts[i], _ = strconv.Atoi(value)
	}
	portsMap := make(map[int]int)
	for _, port := range allPorts {
		for i := 0; i < receiver.ProxyPorts; i++ {
			portsMap[port+i] = 1
		}
	}
	for _, port := range ports {
		delete(portsMap, port)
		value, exist := receiver.listenPorts[port]
		// 如果存在，则关闭本地监听
		if exist {
			value.Close()
			receiver.listenMutex.Lock()
			delete(receiver.listenPorts, port)
			receiver.listenMutex.Unlock()
		}
	}
	for port := range portsMap {
		_, exist := receiver.listenPorts[port]
		if !exist {
			go util.Listen(port, receiver.listenPorts, &receiver.listenMutex, &receiver.allPortTable)
		}
	}
}

func (receiver *Server) serveHttp() {
	// 创建处理请求的处理函数
	queryHandler := func(w http.ResponseWriter, r *http.Request) {
		queryParams := r.URL.Query()
		// 获取特定参数的值
		portParam := queryParams.Get("port")
		port, err := strconv.Atoi(portParam)
		message := ""
		if err == nil {
			host, exist := receiver.allPortTable[port]
			if exist {
				message = host
			}
		}

		// 设置响应头
		w.Header().Set("Content-Type", "text/plain")
		// 写入响应内容
		fmt.Fprintf(w, message)
	}

	listHandler := func(w http.ResponseWriter, r *http.Request) {
		message := ""
		var ports []int
		for port := range receiver.allPortTable {
			ports = append(ports, port)
		}
		sort.Ints(ports)
		for _, port := range ports {
			hostname, _ := receiver.allPortTable[port]
			message += hostname + ":" + strconv.Itoa(port) + "\n"
		}

		// 设置响应头
		w.Header().Set("Content-Type", "text/plain")
		// 写入响应内容
		fmt.Fprintf(w, message)
	}

	// 注册处理函数，并指定路由
	http.HandleFunc("/api/query", queryHandler)
	http.HandleFunc("/api/list", listHandler)

	log.Println("listen http: ", receiver.HttpPort)
	// 启动HTTP服务器，并监听特定的端口
	err := http.ListenAndServe(fmt.Sprintf(":%d", receiver.HttpPort), nil)
	if err != nil {
		log.Fatal("启动HTTP服务器失败:", err)
	}
}
